25 contextlib上下文管理器
with语句你肯定用过——打开文件、获取锁、连接数据库都用它。它能保证资源被正确释放，即使发生异常也不怕。contextlib模块让你更轻松地创建自己的上下文管理器。
一、@contextmanager装饰器
1.1 基本用法
不用写__enter__和__exit__方法，用生成器就能创建上下文管理器。
python
from contextlib import contextmanager
@contextmanager
def managed_resource(name):
print(f"获取资源 {name}")
try:
yield name # 产出值给with语句
finally:
print(f"释放资源 {name}")
# 使用
with managed_resource("数据库") as resource:
print(f"使用 {resource}")输出：
获取资源 数据库
使用 数据库
释放资源 数据库1.2 带异常处理
python
from contextlib import contextmanager
@contextmanager
def error_handler():
try:
yield
except ValueError as e:
print(f"捕获到错误: {e}")
except Exception as e:
print(f"其他错误: {e}")
raise
with error_handler():
raise ValueError("测试错误")
# 输出: 捕获到错误: 测试错误1.3 返回值
python
from contextlib import contextmanager
@contextmanager
def open_file(path, mode):
f = open(path, mode)
try:
yield f # 产出文件对象
finally:
f.close()
with open_file("test.txt", "w") as f:
f.write("Hello")二、实用示例
2.1 计时器
python
from contextlib import contextmanager
import time
@contextmanager
def timer():
start = time.perf_counter()
yield
elapsed = time.perf_counter() - start
print(f"耗时: {elapsed:.4f}秒")
with timer():
sum(range(1000000))
# 输出: 耗时: 0.0312秒2.2 临时修改目录
python
from contextlib import contextmanager
import os
@contextmanager
def change_dir(path):
old_dir = os.getcwd()
os.chdir(path)
try:
yield
finally:
os.chdir(old_dir)
with change_dir("/tmp"):
print(os.getcwd()) # /tmp
print(os.getcwd()) # 原来的目录2.3 临时修改环境变量
python
from contextlib import contextmanager
import os
@contextmanager
def env_var(key, value):
old_value = os.environ.get(key)
os.environ[key] = value
try:
yield
finally:
if old_value is None:
del os.environ[key]
else:
os.environ[key] = old_value
with env_var("DEBUG", "1"):
print(os.environ["DEBUG"]) # 1
print(os.environ.get("DEBUG")) # None2.4 数据库事务
python
from contextlib import contextmanager
@contextmanager
def transaction(connection):
try:
yield connection
connection.commit()
print("事务提交")
except Exception:
connection.rollback()
print("事务回滚")
raise
# 使用
with transaction(conn):
conn.execute("INSERT INTO users ...")三、suppress()：抑制异常
python
from contextlib import suppress
# 文件不存在时不报错
with suppress(FileNotFoundError):
os.remove("nonexistent.txt")
# 等价于
try:
os.remove("nonexistent.txt")
except FileNotFoundError:
pass多个异常：
python
with suppress(FileNotFoundError, PermissionError):
os.remove("somefile.txt")四、redirect_stdout / redirect_stderr
4.1 重定向输出
python
from contextlib import redirect_stdout
import io
# 捕获print输出
f = io.StringIO()
with redirect_stdout(f):
print("这行被重定向了")
print("这行也是")
output = f.getvalue()
print(f"捕获到: {output}")4.2 重定向到文件
python
from contextlib import redirect_stdout
with open("output.txt", "w") as f:
with redirect_stdout(f):
print("写入文件")
print("而不是终端")4.3 同时重定向stdout和stderr
python
from contextlib import redirect_stdout, redirect_stderr
import io
stdout_capture = io.StringIO()
stderr_capture = io.StringIO()
with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
print("stdout内容")
import sys
print("stderr内容", file=sys.stderr)五、closing()：自动关闭
python
from contextlib import closing
# 确保对象的close()方法被调用
with closing(open("file.txt")) as f:
content = f.read()
# 等价于
with open("file.txt") as f:
content = f.read()适用于没有实现上下文管理器协议但有close()方法的对象。
六、nullcontext()：空操作
python
from contextlib import nullcontext
# 有时候不需要上下文管理，但代码结构需要
def process(use_resource=True):
cm = managed_resource("test") if use_resource else nullcontext()
with cm:
print("处理中")
process(True) # 使用资源
process(False) # 不使用资源七、ExitStack：动态管理多个上下文
7.1 基本用法
python
from contextlib import ExitStack
with ExitStack() as stack:
# 动态添加上下文管理器
files = [
stack.enter_context(open(f"file{i}.txt", "w"))
for i in range(3)
]
for f in files:
f.write("hello")
# 所有文件在退出时自动关闭7.2 动态决定上下文
python
from contextlib import ExitStack, suppress
with ExitStack() as stack:
# 根据条件添加不同的上下文
if debug_mode:
stack.enter_context(redirect_stdout(log_file))
stack.enter_context(suppress(FileNotFoundError))
# 正常逻辑
do_something()7.3 注册清理回调
python
from contextlib import ExitStack
def cleanup():
print("清理资源")
with ExitStack() as stack:
stack.callback(cleanup) # 退出时调用
print("正常逻辑")
# 输出:
# 正常逻辑
# 清理资源八、asynccontextmanager：异步版本
python
from contextlib import asynccontextmanager
import asyncio
@asynccontextmanager
async def async_resource():
print("获取异步资源")
try:
yield
finally:
print("释放异步资源")
async def main():
async with async_resource():
print("使用异步资源")
asyncio.run(main())九、实战场景
9.1 日志上下文
python
from contextlib import contextmanager
import logging
@contextmanager
def log_context(logger, level, message):
logger.log(level, f"开始: {message}")
try:
yield
logger.log(level, f"完成: {message}")
except Exception as e:
logger.error(f"失败: {message} - {e}")
raise
logger = logging.getLogger(__name__)
with log_context(logger, logging.INFO, "数据处理"):
process_data()9.2 临时配置
python
from contextlib import contextmanager
@contextmanager
def temp_config(config, **overrides):
original = {k: config[k] for k in overrides}
config.update(overrides)
try:
yield config
finally:
config.update(original)
config = {"debug": False, "verbose": False}
with temp_config(config, debug=True, verbose=True):
print(config) # {'debug': True, 'verbose': True}
print(config) # {'debug': False, 'verbose': False}9.3 重试机制
python
from contextlib import contextmanager
@contextmanager
def retry(max_attempts=3):
for attempt in range(max_attempts):
try:
yield attempt
break
except Exception as e:
if attempt == max_attempts - 1:
raise
print(f"重试 {attempt + 1}/{max_attempts}")
with retry(3) as attempt:
print(f"尝试 {attempt}")
# 可能会抛异常的操作十、总结
contextlib的核心：
| 组件 | 用途 |
|---|---|
@contextmanager | 用生成器创建上下文管理器 |
@asynccontextmanager | 异步版本 |
suppress() | 抑制异常 |
redirect_stdout/stderr | 重定向输出 |
closing() | 自动调用close() |
nullcontext() | 空操作 |
ExitStack | 动态管理多个上下文 |
使用场景：
- 资源管理（文件、数据库连接、网络连接）
- 临时修改状态（环境变量、工作目录、配置）
- 异常处理和清理
- 日志和计时
@contextmanager是最常用的，记住它的用法就够了：yield之前是__enter__，yield之后是__exit__。